package com.cyy.tree;

import java.util.Comparator;

/**
 * 红黑树 红黑树也是一种自平衡的二叉搜索树
 * 以前也叫做平衡二叉B树（Symmetric Binary B Tree）
 * @author 会玩的洋洋
 * @since 2022-02-26
 * @param <E>
 */
public class RBTree<E> extends BBST<E> {
    /**
     * 定义 红色节点 为 false
     */
    private static final boolean RED = false;
    /**
     * 定义 黑色节点 为 true
     */
    private static final boolean BLACK = true;

    public RBTree() {
    }

    public RBTree(Comparator<E> comparator) {
        super(comparator);
    }

    /**
     * 添加之后进行调整红黑树性质
     * @param node 添加的节点
     */
    @Override
    protected void afterAdd(BinaryTree.Node<E> node) {
        if (node == null) {
            return;
        }
        // 获取父节点
        BinaryTree.Node<E> parent = node.parent;
        if (parent == null) {
            // 添加的根节点
            black(node);
            return;
        }
        // parent 为 BLACK 满足性质4
        if (colorOf(parent)) {
            return;
        }
        // parent 为 RED 分8种情况
        // 获取叔父节点
        BinaryTree.Node<E> uncle = parent.sibling();
        // 获取祖父节点
        BinaryTree.Node<E> grand = red(parent.parent);
        // 先判断 上溢情况 叔父节点为 RED
        if (isRed(uncle)) {
            // 将叔父节点和父节点染成黑色  祖父节点染成红色
            black(uncle);
            black(parent);
            // 把祖父节点当做新添加的节点
            afterAdd(grand);
            return;
        }

        // 上溢情况  叔父节点为 BLACK （相当于没有叔父节点）
        if (parent.isLeftChildren()) {
            if (node.isLeftChildren()) {
                // LL
                black(parent);
            } else {
                // LR
                black(node);
                leftRotate(parent);
            }
            rightRotate(grand);
        } else {
            if (node.isRightChildren()) {
                // RR
                black(parent);
            } else {
                // RL
                black(node);
                rightRotate(parent);
            }
            leftRotate(grand);
        }
    }

    /**
     * 删除之后调整红黑树性质
     * @param node 删除的节点
     */
    @Override
    protected void afterRemove(BinaryTree.Node<E> node) {
        // 如果删除的是 RED 节点返回
        // 如果要删除的节点是一个 RED 节点的 BLACK 节点，即替代的节点是红色的
        if (isRed(node)) {
            black(node);
            return;
        }

        BinaryTree.Node<E> parent = node.parent;
        // 如果删除的是根节点
        if (parent == null) {
            return;
        }

        // 删除的是黑色叶子节点【下溢】
        // 判断被删除的节点是左还是右
        boolean left = parent.left == null || node.isLeftChildren();
        BinaryTree.Node<E> sibling = left ? parent.right : parent.left;

        // 删除黑色叶子节点 情况
        // 删除的节点是左边的情况
        if (left) {
            // 兄弟节点为 红色 的情况
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                leftRotate(parent);
                sibling = parent.right;
            }

            // 兄弟节点必然为黑色
            // 兄弟节点没有 RED 子节点 情况
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean parentBlack = isBlack(parent);
                red(sibling);
                black(parent);
                if (parentBlack) {
                    // parent 可能下溢
                    afterRemove(parent);
                }
            } else {
                // 兄弟节点存在 RED 子节点,直接向兄弟节点借一个节点
                if (isBlack(sibling.right)) {
                    // 如果兄弟节点的右边为null 即 black 则先进行右旋转,将子节点放到 silbing.right中
                    rightRotate(sibling);
                    sibling = parent.right;
                }
                // 染色
                color(sibling, colorOf(parent));
                black(parent);
                black(sibling.right);
                // 进行左旋转
                leftRotate(parent);
            }
        } else {
            // 删除的是右节点 与 删除左节点步骤相同，方向相反
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rightRotate(parent);
                sibling = parent.left;
            }

            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean parentBlack = isBlack(parent);
                red(sibling);
                black(parent);
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else {
                if (isBlack(sibling.left)) {
                    leftRotate(sibling);
                    sibling = parent.left;
                }
                color(sibling, colorOf(parent));
                black(parent);
                black(sibling.left);
                rightRotate(parent);
            }
        }

    }

    /**
     * 创建红黑树节点
     * @param element 元素
     * @param parent 父节点
     * @return Node<E>
     */
    @Override
    protected BinaryTree.Node<E> createNode(E element, BinaryTree.Node<E> parent) {
        return new RBNode<>(parent, element);
    }

    /**
     * 染色操作 将特定的 node 染色成为指定的 color
     * @param node 特定的node
     * @param color 指定的颜色
     * @return Node<E>
     */
    private BinaryTree.Node<E> color(BinaryTree.Node<E> node, boolean color) {
        if (node == null) {
            return null;
        }
        ((RBNode<E>) node).color = color;
        return node;
    }

    /**
     * 染色 将节点染成红色
     * @param node 需要染色的节点
     * @return Node<E>
     */
    private BinaryTree.Node<E> red(BinaryTree.Node<E> node) {
        return color(node, RED);
    }

    /**
     * 染色 将节点染成黑色
     * @param node 需要染色的节点
     * @return Node<E>
     */
    private BinaryTree.Node<E> black(BinaryTree.Node<E> node) {
        return color(node, BLACK);
    }

    /**
     * 判断节点是什么颜色
     * @param node 需要判断的节点
     * @return boolean 节点颜色
     */
    private boolean colorOf(BinaryTree.Node<E> node) {
        // 如果节点为 null 则为虚拟叶子节点 即为 黑色
        return node == null ? BLACK : ((RBNode<E>) node).color;
    }

    /**
     * 判断节点是否为黑色
     * @param node 需要判断的节点
     * @return boolean
     */
    private boolean isBlack(BinaryTree.Node<E> node) {
        return node == null || ((RBNode<E>) node).color == BLACK;
    }

    /**
     * 判断节点是否为红色
     * @param node 需要判断的节点
     * @return boolean
     */
    private boolean isRed(BinaryTree.Node<E> node) {
        return colorOf(node) == RED;
    }

    private static class RBNode<E> extends BinaryTree.Node<E> {
        // 默认为 RED，可以解决一系列很多的红黑树性质
        boolean color = RED;

        public RBNode(BinaryTree.Node<E> parent, E element) {
            super(parent, element);
        }

        @Override
        public String toString() {
            String str = "";
            if (color == RED) {
                str = "RED_";
            } else {
                str = "BLACK_";
            }
            return str + element;
        }
    }
}
